#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Mar 31 16:22:52 2024

@author: lixiao
"""
import numpy as np
from typing import Callable

from perpectiveGeometry import (
    SE3_to_se3,
    se3_to_SE3
  )
from utils import unvectorizeLieParameters

def reprojection_error(Ps, X, images):
  '''
  Ps: projection matrices (M x 3 x 4)
  X: 3D points in homogeneous coordinates (4 x N)
  images - (M x 3 x N), 
    We assume that all feature points are viewed by all cameras
  '''
  # Number of views
  num_views = Ps.shape[0]
  # Number of (3D) points
  num_points = X.shape[1]
  
  # Initialize the reprojection error vector
  reproj_err = np.zeros(2 * num_views * num_points)
  
  # Compute reprojection errors
  for i in range(num_views):
    # Project 3D points to 2D
    x_hat = np.dot(Ps[i], X)
    x_hat = x_hat / x_hat[2, :]
    
    # Compute reprojection error with respect to the observed 2D points
    err_temp = x_hat[0:2, :] - images[i][0:2, :]
    err_temp = err_temp.T.reshape(-1)
    # Set small errors to zeros
    err_temp[np.where(np.abs(err_temp) <= 1e-5)] = 0.0
    reproj_err[i*2*num_points : (i+1)*2*num_points] = err_temp
  
  return reproj_err

def lieJacobian(para_vec, K, M, N):
  '''
  Computes the Jacobian of the reprojection error with respect to 
  the motion parameters (se(3) vectors)
  and the structure parameters (3D points).

  M - Number of views
  N - Number of features
  '''

  fx = K[0,0]
  fy = K[1,1]
  J = np.zeros((2*N*M, (6*M + 3*N)))
  
  for i in range(M):
    T = se3_to_SE3(para_vec[i*6:(i+1)*6])
    
    for j in range(N):
      Xp = para_vec[6*M+j*3:6*M+(j+1)*3]
      norm_xp = T @ np.append(Xp, 1)
      x = norm_xp[0]
      y = norm_xp[1]
      z = norm_xp[2]
      
      J_motion = np.array([
        [fx * 1/z, 0, -fx * (x/z**2), -fx * (x * y / z**2), fx * (1 + x**2/z**2), -fx * y/z],
        [0, fy * 1/z, -fy * (y/z**2), -fy * (1 + y**2/z**2), fy * x * y / z**2, fy * x/z]
        ])
      
      J_point = np.array([
        [fx / z, 0, -fx * x / z**2],
        [0, fy / z, -fy * y / z**2]
        ]) @ T[0:3,0:3]
      
      J[2*N*i+2*j : 2*N*i+2*(j+1), 6*i : 6*(i+1)] = J_motion
      J[2*N*i+2*j : 2*N*i+2*(j+1),6*M+3*j : 6*M+3*(j+1)] = J_point
      
  return -J



def lm(paramVec, img_pts, K, M, N, fx:Callable, Jx:Callable, k_max = 100,
       tolerance = 1e-16):
  '''
  Args:
    paramVec - ((6*M + 3*N) x 1)
    img_pts - (M x 3 x N), 
      We assume that all feature points are viewed by all cameras
    K - Camera intrisinc matrix
    M - Number of views
    N - Number of features
  '''
  stop = False
  norm = np.linalg.norm
  solve = np.linalg.solve
  
  J = Jx(paramVec, K, M, N)
  residual = fx(paramVec, img_pts, K, M, N)
  A = J.T@J
  g = J.T@residual
  dampingCoeff = 1e-1 * np.max(np.diag(A))
  err = norm(residual, 2)
  
  for i in range(k_max):
    if stop:
      print('stop!')
      break
    
    deltap = solve(A+dampingCoeff*np.eye(A.shape[0]), g)
    
    # if norm(deltap, np.inf) < tolerance*norm(paramVec,np.inf):
    #   stop = True
    
    paramVec_new = paramVec - deltap
    residual_new = fx(paramVec_new, img_pts, K, M, N)
    err_new = norm(residual_new, 2)
    # print(err_new < err)
    
    if err_new < err:
      if np.fabs(err_new - err) < tolerance:
        stop = True
      else:
        paramVec = paramVec_new
        residual = residual_new
        J = Jx(paramVec, K, M, N)
        g = J.T@residual
        A = J.T@J
        err = norm(residual, 2)
        stop = (norm(g, 2) <= tolerance) or (err <= tolerance)
        dampingCoeff = dampingCoeff / 2.
    else:
      dampingCoeff = dampingCoeff * 2.
  
  Ps_opt, X_opt = unvectorizeLieParameters(paramVec, K, M, N)
  err_opt = err
  return Ps_opt, X_opt, err_opt


def gn(paramVec, img_pts, K, M, N, fx:Callable=None, Jx:Callable = lieJacobian, 
       k_max = 200, tolerance = 1e-16):
  
  def fx_(para_vec, x, K, M, N):
    Ps, X = unvectorizeLieParameters(para_vec, K, M, N)
    return reprojection_error(Ps, X, x)
  
  if fx is None:
    fx = fx_
  
  norm = np.linalg.norm
  solve = np.linalg.solve
  lr = 1e-2
  
  J = Jx(paramVec, K, M, N)
  residual = fx(paramVec, img_pts, K, M, N)
  A = J.T@J
  g = J.T@residual
  err = norm(residual, 2)
  
  for i in range(k_max):
    deltap = solve(A+np.eye(A.shape[0]), g)
    paramVec = paramVec + (lr*deltap)
        
    residual = fx(paramVec, img_pts, K, M, N)
    err = norm(residual,2)
    # print(err)

    J = Jx(paramVec, K, M, N)
    g = J.T@residual
    A = J.T@J
  
  Ps_opt, X_opt = unvectorizeLieParameters(paramVec, K, M, N)
  return Ps_opt, X_opt, err


def simple_lm(paramVec, img_pts, K, M, N, fx:Callable=None, Jx:Callable = lieJacobian, 
       k_max = 500, tolerance = 1e-16):
  '''
  A rudimentary implementation of the Levenberg-Marquardt (LM) algorithm.
  '''
  def fx_(para_vec, x, K, M, N):
    Ps, X = unvectorizeLieParameters(para_vec, K, M, N)
    return reprojection_error(Ps, X, x)
  
  if fx is None:
    fx = fx_
  
  norm = np.linalg.norm
  solve = np.linalg.solve
  lr = 1e-2
  
  J = Jx(paramVec, K, M, N)
  residual = fx(paramVec, img_pts, K, M, N)
  A = J.T@J
  g = J.T@residual
  err = norm(residual, 2)
  dampingCoeff = 1e-1 * np.max(np.diag(A))
  t = 0
  for i in range(k_max):
    deltap = solve(A+dampingCoeff*np.eye(A.shape[0]), g)
    paramVec_new = paramVec + deltap
        
    residual = fx(paramVec_new, img_pts, K, M, N)
    err_new = norm(residual,np.inf)
    # print(err - err_new)
    if err_new < err:
      paramVec = paramVec_new
      residual = fx(paramVec, img_pts, K, M, N)
      err = norm(residual,np.inf)
      J = Jx(paramVec, K, M, N)
      g = J.T@residual
      A = J.T@J
      dampingCoeff = dampingCoeff / 2.
      t += 1
      r = residual
      # print(err)
      # print('coeff', dampingCoeff)
    else:
      dampingCoeff = dampingCoeff * 2.
  err = norm(r,2)
  # print(t)
  Ps_opt, X_opt = unvectorizeLieParameters(paramVec, K, M, N)
  return Ps_opt, X_opt, err



